import warnings
warnings.filterwarnings('ignore')
import argparse, yaml, copy
import numpy as np
from ultralytics.models.yolo.detect.compress import DetectionCompressor, DetectionFinetune
from ultralytics.models.yolo.segment.compress import SegmentationCompressor, SegmentationFinetune
from ultralytics.models.yolo.pose.compress import PoseCompressor, PoseFinetune
from ultralytics.models.yolo.obb.compress import OBBCompressor, OBBFinetune

def compress(param_dict):
    with open(param_dict['sl_hyp'], errors='ignore') as f:
        sl_hyp = yaml.safe_load(f)
    param_dict.update(sl_hyp)
    # 剔除自定义参数
    attn_weight_path = param_dict.pop("attn_weight_path", None)
    attn_alpha = param_dict.pop("attn_alpha", 0.7)
    param_dict['name'] = f'{param_dict["name"]}-prune'
    param_dict['patience'] = 0
    compressor = DetectionCompressor(overrides=param_dict)
    compressor.attn_weight_path = attn_weight_path
    compressor.attn_alpha = attn_alpha
    # print(f"1111111attn_weight_path: {attn_weight_path}")
    # print(f"11111attn_alpha: {attn_alpha}" ) 
    # print(f"11111111compressor: {compressor}")
    # compressor = SegmentationCompressor(overrides=param_dict)
    # compressor = PoseCompressor(overrides=param_dict)
    # compressor = OBBCompressor(overrides=param_dict)
    prune_model_path = compressor.compress()
    return prune_model_path

def finetune(param_dict, prune_model_path):
    param_dict['model'] = prune_model_path
    param_dict['name'] = f'{param_dict["name"]}-finetune'
    trainer = DetectionFinetune(overrides=param_dict)
    # trainer = SegmentationFinetune(overrides=param_dict)
    # trainer = PoseFinetune(overrides=param_dict)
    # trainer = OBBFinetune(overrides=param_dict)
    trainer.train()

if __name__ == '__main__':
    param_dict = {
        # origin
        'model': r'runs\train\yolov8-c2famhabi\weights\best_v2.pt',
        'data':'../dataset/data.yaml',
        'imgsz': 640,
        'epochs': 300,
        'batch': 1,
        'workers': 0,
        'cache': False,
        'optimizer': 'SGD',
        'device': '0',
        'close_mosaic': 0,
        'project':'runs/prune',
        'name':'yolov8n-visdrone-lamp-exp',
        
        # prune
        'prune_method':'group_taylor',
        'global_pruning': True,
        'speed_up': 3.0,
        'reg': 0.0005,
        'sl_epochs': 500,
        'sl_hyp': 'ultralytics/cfg/hyp.scratch.sl.yaml',
        'sl_model': None,
    }
    
    prune_model_path = compress(copy.deepcopy(param_dict))
    finetune(copy.deepcopy(param_dict), prune_model_path)